'''
@author: zhangkai
@license: (C) Copyright 2017-2023
@contact: jeffcobile@gmail.com
@Software : PyCharm
@file: base.py
@time: 2020-09-14 22:52:41
@desc: 
'''
import os
from jjzhk.config import ZKCFG
import tqdm
import cv2
import xml.etree.ElementTree as ET
import pycocotools.coco as coco
from pycocotools import mask as coco_mask
import numpy as np


class DataSetBase(object):
    def __init__(self, root,imageset):
        self._root_path = root
        self._imageset = imageset
        self._img_list = []
        self.__prepare__()

    def __len__(self):
        return len(self._img_list)

    def __getitem__(self, index):
        image = self.getImage(index)
        mask = self.getMask(index)
        return image, mask

    def __prepare__(self):
        self.prepare()

    def __getItemInfo__(self, index):
        return self._img_list[index]

    def __getItemInfoByImageId__(self, imgid):
        pass

    def __getIndexByImageId__(self, imgid):
        index = 0
        for item in self._img_list:
            if item["img_id"] == imgid:
                return index

            index += 1

        return -1

    def prepare(self):
        pass

    def getImage(self, index):
        pass

    def getMask(self, index):
        pass


class VOCDataSet(DataSetBase):
    def __init__(self, config:ZKCFG, imageset):
        root = os.path.join(config.BASE.DATA_ROOT, config.BASE.DATA_SUB_ROOT)
        self.dataConfig = config
        super(VOCDataSet, self).__init__(root=root, imageset=imageset)

    def prepare(self):
        prefix = "seg"
        image_list_file = os.path.join(self._root_path, "MainSet", "%s_%s.txt" % (prefix, self._imageset))
        lines = [x.strip() for x in open(image_list_file, 'r').readlines()]

        bar = tqdm.tqdm(lines)
        for line in bar:
            bar.set_description("Processing %s" % self._imageset)
            l = self.__getItemInfoByImageId__(line)
            self._img_list.append(l)

    def getImage(self, index):
        fname = os.path.join(self._root_path, "JPEGImages", "%s.jpg" % self._img_list[index]["img_id"])
        self._img_list[index]["path"] = fname
        return fname

    def getMask(self, index):
        info = self._img_list[index]
        file_index = info["img_id"]
        lbl_path = os.path.join(self._root_path, 'SegmentationDecode', file_index + '.png')
        # label_mask = cv2.imread(lbl_path, cv2.IMREAD_GRAYSCALE)

        return lbl_path

    '''
    {
        img_id : "***",
        width : "***",
        height : "***",
        boxes : [[classs_id, xmin, ymin, xmax, ymax],[classs_id, xmin, ymin, xmax, ymax]
        detail : [{"name":***,"pose":***, "truncated":***,"difficult":***,},{....},{....}]
        ...]
    }
    '''
    def __getItemInfoByImageId__(self, image_id):
        image_info = {}
        image_info["img_id"] = image_id

        anno_file = os.path.join(self._root_path, "Annotations", "%s.xml" % image_id)
        root = ET.parse(anno_file).getroot()
        size = root.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)
        image_info["width"] = w
        image_info["height"] = h

        # labels = []
        # detail = []
        # for obj in root.iter('object'):
        #     class_name = obj.find('name').text
        #     xmlbox = obj.find('bndbox')
        #     left   = int(xmlbox.find('xmin').text)
        #     right  = int(xmlbox.find('xmax').text)
        #     top    = int(xmlbox.find('ymin').text)
        #     bottom = int(xmlbox.find('ymax').text)
        #     class_id = self.dataConfig.classid(class_name)
        #     item = [class_id, left, top, right, bottom]
        #     labels.append(item)
        #
        #     pose = obj.find("pose").text
        #     truncated = int(obj.find("truncated").text)
        #     difficult = int(obj.find("difficult").text)
        #     detail.append(
        #         {"name": class_name,
        #          "pose":pose,
        #          "truncated" : truncated,
        #          "difficult" : difficult,
        #          "bbox" : [left, top, right, bottom]
        #          })
        #
        # image_info["boxes"] = labels
        # image_info["detail"] = detail

        return image_info


class COCODataSet(DataSetBase):
    def __init__(self, config:ZKCFG, imageset):
        self.dataConfig = config
        self.year = 2017
        imageset = imageset.replace('eval', 'val')
        self.prefix = "%s%d" % (imageset, self.year)
        self.CAT_LIST = [0, 5, 2, 16, 9, 44, 6, 3, 17, 62, 21, 67, 18, 19, 4,
                1, 64, 20, 63, 7, 72]
        super(COCODataSet, self).__init__(root=os.path.join(config.BASE.DATA_ROOT, config.BASE.DATA_SUB_ROOT), imageset=imageset)

    def prepare(self):
        self.annFile = os.path.join(self._root_path, "annotations", "instances_%s.json" % self.prefix)
        print(self.annFile)
        self.coco = coco.COCO(self.annFile)

        # class_ids = sorted(self.coco.getCatIds())  # all classes
        # lines = []
        # for id in class_ids:
        #     lines.extend(list(self.coco.getImgIds(catIds=[id])))
        ids = list(self.coco.imgs.keys())
        lines = self._preprocess_(ids)

        lines = list(set(lines))

        bar = tqdm.tqdm(lines)
        for line in bar:
            bar.set_description("Processing %s" % self._imageset)
            l = self.__getItemInfoByImageId__(line)
            self._img_list.append(l)

    def getImage(self, index):
        info = self._img_list[index]
        return os.path.join(self._root_path, self.prefix, info['path'])

    def getMask(self, index):
        img_id = self._img_list[index]['img_id']
        img_metadata = self.coco.loadImgs(img_id)[0]
        cocotarget = self.coco.loadAnns(self.coco.getAnnIds(imgIds=img_id))
        # mask = Image.fromarray(self._gen_seg_mask_s(
        #     cocotarget, img_metadata['height'], img_metadata['width']))

        return (cocotarget, img_metadata)

        # info = self._img_list[index]
        # img_id = info['img_id']
        # height = info['height']
        # width = info['width']
        # annIds = self.coco.getAnnIds(imgIds=img_id, iscrowd=None)
        # anns = self.coco.loadAnns(annIds)
        # MASK = np.zeros((height, width), dtype=np.uint8)
        #
        # for i, ann in enumerate(anns):
        #     mask = self.coco.annToMask(ann)
        #     idxs = np.where(mask > 0)
        #     MASK[idxs] = ann['category_id']
        #
        # masks = [self.coco.annToMask(obj).reshape(-1) for obj in anns]
        # masks = np.vstack(masks)
        # masks = masks.reshape(-1, info["height"], info["width"])
        # self._img_list[index]["mask"] = masks
        #
        # return MASK

    '''
    {
        img_id : "***",
        width : "***",
        height : "***",
        boxes : [[classs_id, xmin, ymin, xmax, ymax],[classs_id, xmin, ymin, xmax, ymax]
        ...]
    }
    '''
    def __getItemInfoByImageId__(self, image_id):
        image_info = self.coco.imgs[image_id]
        img_info_json = {}
        img_info_json["img_id"] = image_id
        img_info_json["width"] = image_info["width"]
        img_info_json["height"] = image_info["height"]

        annIds = self.coco.getAnnIds(imgIds=image_id, iscrowd=None)
        ann_info = self.coco.loadAnns(annIds)
        img_info_json["anno"] = ann_info

        img_info = self.coco.imgs[image_id]
        img_info_json["path"] = img_info['file_name']

        return img_info_json

    def _preprocess_(self, ids):
        new_ids = []
        for img_id in ids:
            cocotarget = self.coco.loadAnns(self.coco.getAnnIds(imgIds=img_id))
            img_metadata = self.coco.loadImgs(img_id)[0]
            mask = self._gen_seg_mask_(cocotarget, img_metadata['height'], img_metadata['width'])
            # more than 1k pixels
            if (mask > 0).sum() > 1000:
                new_ids.append(img_id)

        return new_ids

    def _gen_seg_mask_(self, target, h, w):
        mask = np.zeros((h, w), dtype=np.uint8)
        for instance in target:
            rle = coco_mask.frPyObjects(instance['segmentation'], h, w)
            m = coco_mask.decode(rle)
            cat = instance['category_id']
            if cat in self.CAT_LIST:
                c = self.CAT_LIST.index(cat)
            else:
                continue
            if len(m.shape) < 3:
                mask[:, :] += (mask == 0) * (m * c)
            else:
                mask[:, :] += (mask == 0) * (((np.sum(m, axis=2)) > 0) * c).astype(np.uint8)
        return mask


class TestImageDataSet(DataSetBase):
    def __init__(self, cfg:ZKCFG):
        self.cfg = cfg
        super(TestImageDataSet, self).__init__(root=cfg.BASE.TEST_DATA_ROOT, imageset='test')

    def prepare(self):
        image_list_file = os.path.join(self._root_path, "info.txt")
        lines = [x.strip() for x in open(image_list_file, 'r').readlines()]

        bar = tqdm.tqdm(lines)
        for line in bar:
            bar.set_description("Processing %s" % self._imageset)
            self._img_list.append({"img_id" : line})

    def getImage(self, index):
        fname = os.path.join(self._root_path, "Images", "%s.jpg" % self._img_list[index]["img_id"])
        image = cv2.imread(fname)
        self._img_list[index]["path"] = fname
        self._img_list[index]["height"] = image.shape[0]
        self._img_list[index]["width"] = image.shape[1]
        return fname

    def getMask(self, index):
        return None


class ADE20KDataSet(DataSetBase):
    def __init__(self, config:ZKCFG, imageset, isSegmentation=True):
        root = os.path.join(config.BASE.DATA_ROOT, config.BASE.DATA_SUB_ROOT)
        isSegmentation = True # must be true
        self.dataConfig = config
        super(ADE20KDataSet, self).__init__(root=root, imageset=imageset, isSegmentation=isSegmentation)

    def prepare(self):
        if self._imageset == "train":
            lines = self.__recursive_glob__(os.path.join(self._root_path, "training"), ".jpg")
        else:
            lines = self.__recursive_glob__(os.path.join(self._root_path, "validation"), ".jpg")

        bar = tqdm.tqdm(lines)
        for line in bar:
            bar.set_description("Processing %s" % self._imageset)
            l = self.__getItemInfoByImageId__(os.path.basename(line))
            l["path"] = line
            self._img_list.append(l)

    def getImage(self, index):
        fname = self._img_list[index]["path"]
        image = cv2.imread(fname)

        self._img_list[index]["width"] = image.shape[1]
        self._img_list[index]["height"] = image.shape[0]

        return image

    def getLabel(self, index):
        return self._img_list[index]["boxes"]

    def getMask(self, index):
        lbl_path = self._img_list[index]["path"][:4] + "_seg.png"
        label_mask = cv2.imread(lbl_path)

        return label_mask

    '''
    {
        img_id : "***",
        width : "***",
        height : "***",
        boxes : [[classs_id, xmin, ymin, xmax, ymax],[classs_id, xmin, ymin, xmax, ymax]
        detail : [{"name":***,"pose":***, "truncated":***,"difficult":***,},{....},{....}]
        ...]
    }
    '''
    def __getItemInfoByImageId__(self, image_id):
        image_info = {}
        image_info["img_id"] = image_id

        # image_info["width"] = w
        # image_info["height"] = h

        image_info["boxes"] = []

        return image_info

    def __recursive_glob__(self, rootdir=".", suffix=""):
        """Performs recursive glob with given suffix and rootdir
            :param rootdir is the root directory
            :param suffix is the suffix to be searched
        """
        return [
            os.path.join(looproot, filename)
            for looproot, _, filenames in os.walk(rootdir)
            for filename in filenames
            if filename.endswith(suffix)
        ]


class PennFudanDataSet(DataSetBase):
    def __init__(self, config: ZKCFG, imageset, isSegmentation):
        root = os.path.join(config.BASE.DATA_ROOT, config.BASE.DATA_SUB_ROOT)
        self.dataConfig = config
        super(PennFudanDataSet, self).__init__(root=root, imageset=imageset, isSegmentation=isSegmentation)

    def prepare(self):
        lines = list(sorted(os.listdir(os.path.join(self._root_path, "PNGImages"))))

        bar = tqdm.tqdm(lines)
        for line in bar:
            bar.set_description("Processing %s" % self._imageset)
            line = line.replace(".png", "")
            l = self.__getItemInfoByImageId__(line)
            self._img_list.append(l)

    def __getItemInfoByImageId__(self, image_id):
        image_info = {}
        image_info["img_id"] = image_id
        mask_path = os.path.join(self._root_path, "PedMasks", "%s_mask.png" % image_id)
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        image_info["width"] = mask.shape[1]
        image_info["height"] = mask.shape[0]

        obj_ids = np.unique(mask)
        obj_ids = obj_ids[1:] # 第一个是背景
        masks = mask == obj_ids[:, None, None]
        num_objs = len(obj_ids)
        boxes = []
        for i in range(num_objs):
            pos = np.where(masks[i])
            xmin = np.min(pos[1])
            xmax = np.max(pos[1])
            ymin = np.min(pos[0])
            ymax = np.max(pos[0])
            boxes.append([xmin, ymin, xmax, ymax])

        boxes = np.array(boxes, dtype=np.float32)
        labels = np.ones((num_objs,), dtype=np.float32)
        boxes = np.hstack((labels[:, np.newaxis], boxes)).astype(np.float32, copy=False)

        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        iscrowd = np.zeros((num_objs,), dtype=np.int64)

        image_info["boxes"] = boxes
        image_info["labels"] = labels
        image_info["masks"] = masks
        image_info["area"] = area
        image_info["iscrowd"] = iscrowd
        return image_info

    def getImage(self, index):
        fname = os.path.join(self._root_path, "PNGImages", "%s.png" % self._img_list[index]["img_id"])
        self._img_list[index]["path"] = fname
        return cv2.imread(fname)

    def getLabel(self, index):
        return self._img_list[index]["boxes"]

    def getMask(self, index):
        return self._img_list[index]["masks"]


class CityScapesDataSet(DataSetBase):
    def __init__(self, config:ZKCFG, imageset):
        root = os.path.join(config.BASE.DATA_ROOT, config.BASE.DATA_SUB_ROOT)
        self.dataConfig = config
        super(CityScapesDataSet, self).__init__(root=root, imageset=imageset)

    def prepare(self):
        lines = []
        img_folder = os.path.join(self._root_path, 'leftImg8bit/' + self._imageset)

        for root, _, files in os.walk(img_folder):
            for filename in files:
                if filename.endswith('.png'):
                    imgpath = os.path.join(root, filename)
                    lines.append(imgpath)
        bar = tqdm.tqdm(lines)
        for line in bar:
            bar.set_description("Processing %s" % self._imageset)
            l = self.__getItemInfoByImageId__(os.path.basename(line))
            l["path"] = line
            self._img_list.append(l)

    def getImage(self, index):
        fname = self._img_list[index]["path"]
        # print(fname)
        image = cv2.imread(fname)

        self._img_list[index]["width"] = image.shape[1]
        self._img_list[index]["height"] = image.shape[0]

        return fname

    def getMask(self, index):
        fname = self._img_list[index]["path"]
        lbl_path = fname.replace('leftImg8bit', 'gtFine')
        lbl_path = "%s_%s.png" % (lbl_path[:-4], "labelIds")
        # label_mask = cv2.imread(lbl_path, cv2.IMREAD_GRAYSCALE)

        return lbl_path

    '''
    {
        img_id : "***",
        width : "***",
        height : "***",
        boxes : [[classs_id, xmin, ymin, xmax, ymax],[classs_id, xmin, ymin, xmax, ymax]
        detail : [{"name":***,"pose":***, "truncated":***,"difficult":***,},{....},{....}]
        ...]
    }
    '''
    def __getItemInfoByImageId__(self, image_id):
        image_info = {}
        image_info["img_id"] = image_id

        # image_info["width"] = w
        # image_info["height"] = h

        image_info["boxes"] = []

        return image_info




